{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import numpy as np\n",
    "from scipy.optimize import curve_fit\n",
    "from scipy import asarray as ar,exp\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib notebook\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "infect_A=pd.read_csv('infectDataset/city_A/infection.csv',header=None)\n",
    "infect_A.columns = ['city','region_id','date', 'index']\n",
    "infect_B=pd.read_csv('infectDataset/city_B/infection.csv',header=None)\n",
    "infect_B.columns = ['city','region_id','date', 'index']\n",
    "infect_C=pd.read_csv('infectDataset/city_C/infection.csv',header=None)\n",
    "infect_C.columns = ['city','region_id','date', 'index']\n",
    "infect_D=pd.read_csv('infectDataset/city_D/infection.csv',header=None)\n",
    "infect_D.columns = ['city','region_id','date', 'index']\n",
    "infect_E=pd.read_csv('infectDataset/city_E/infection.csv',header=None)\n",
    "infect_E.columns = ['city','region_id','date', 'index']\n",
    "infect_F=pd.read_csv('infectDataset/city_F/infection.csv',header=None)\n",
    "infect_F.columns = ['city','region_id','date', 'index']\n",
    "infect_G=pd.read_csv('infectDataset/city_G/infection.csv',header=None)\n",
    "infect_G.columns = ['city','region_id','date', 'index']\n",
    "infect_H=pd.read_csv('infectDataset/city_H/infection.csv',header=None)\n",
    "infect_H.columns = ['city','region_id','date', 'index']\n",
    "infect_I=pd.read_csv('infectDataset/city_I/infection.csv',header=None)\n",
    "infect_I.columns = ['city','region_id','date', 'index']\n",
    "infect_J=pd.read_csv('infectDataset/city_J/infection.csv',header=None)\n",
    "infect_J.columns = ['city','region_id','date', 'index']\n",
    "infect_K=pd.read_csv('infectDataset/city_K/infection.csv',header=None)\n",
    "infect_K.columns = ['city','region_id','date', 'index']\n",
    "#\n",
    "#\n",
    "infection_data={'A':infect_A,'B':infect_B,'C':infect_C,'D':infect_D,'E':infect_E,'F':infect_F,'G':infect_G,'H':infect_H,'I':infect_I,'J':infect_J,'K':infect_K}\n",
    "#按照date进行groupby,total_A就表示了city_A 45天每天的新增人数\n",
    "grouped=infect_A.groupby(['date'])\n",
    "total_A=grouped['index'].agg(np.sum).reset_index(drop=True)#45天每天新增人数\n",
    "#acc_A=[np.sum(total_A[:i]) for i in range(1,len(total_A)+1)]#45天每天累计感染人数\n",
    "grouped=infect_B.groupby(['date'])\n",
    "total_B=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_B=[np.sum(total_B[:i]) for i in range(1,len(total_B)+1)]\n",
    "grouped=infect_C.groupby(['date'])\n",
    "total_C=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_C=[np.sum(total_C[:i]) for i in range(1,len(total_C)+1)]\n",
    "grouped=infect_D.groupby(['date'])\n",
    "total_D=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_D=[np.sum(total_D[:i]) for i in range(1,len(total_D)+1)]\n",
    "grouped=infect_E.groupby(['date'])\n",
    "total_E=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_E=[np.sum(total_E[:i]) for i in range(1,len(total_E)+1)]\n",
    "grouped=infect_F.groupby(['date'])\n",
    "total_F=grouped['index'].agg(np.sum).reset_index(drop=True)#45天每天新增人数\n",
    "#acc_F=[np.sum(total_F[:i]) for i in range(1,len(total_F)+1)]#45天每天累计感染人数\n",
    "grouped=infect_G.groupby(['date'])\n",
    "total_G=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_G=[np.sum(total_G[:i]) for i in range(1,len(total_G)+1)]\n",
    "grouped=infect_H.groupby(['date'])\n",
    "total_H=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_H=[np.sum(total_H[:i]) for i in range(1,len(total_H)+1)]\n",
    "grouped=infect_I.groupby(['date'])\n",
    "total_I=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_I=[np.sum(total_I[:i]) for i in range(1,len(total_I)+1)]\n",
    "grouped=infect_J.groupby(['date'])\n",
    "total_J=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_J=[np.sum(total_J[:i]) for i in range(1,len(total_J)+1)]\n",
    "grouped=infect_K.groupby(['date'])\n",
    "total_K=grouped['index'].agg(np.sum).reset_index(drop=True)\n",
    "#acc_K=[np.sum(total_K[:i]) for i in range(1,len(total_K)+1)]\n",
    "city_infect={'A':total_A,'B':total_B,'C':total_C,'D':total_D,'E':total_E,'F':total_F,'G':total_G,'H':total_H,'I':total_I,'J':total_J,'K':total_K}\n",
    "#acc_infect={'A':acc_A,'B':acc_B,'C':acc_C,'D':acc_D,'E':acc_E,'F':acc_F,'G':acc_G,'H':acc_H,'I':acc_I,'J':acc_J,'K':acc_K}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nre_sub=pd.read_csv('submits/submissionShilft1_16.csv',header=None)\\nre_sub.columns = ['city','region_id','date', 'index']\\npre_city_id={}\\npre_id={}\\npre_best={}\\ndef get_true(city,Id):\\n    infection=infection_data[city]\\n    id_infect=infection[infection['region_id'].isin([Id])].reset_index(drop=True)\\n    return id_infect['index'].values\\ngrouped=re_sub.groupby(['city'])\\nfor city,group in grouped:\\n    data=group.reset_index(drop=True)\\n    grouped1=data.groupby(['date'])\\n    pre_city=grouped1['index'].agg(np.sum).reset_index(drop=True)#30天每天新增人数\\n    pre_best[city]=pre_city\\n\""
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#这段代码\n",
    "'''\n",
    "re_sub=pd.read_csv('submits/submissionShilft1_16.csv',header=None)\n",
    "re_sub.columns = ['city','region_id','date', 'index']\n",
    "pre_city_id={}\n",
    "pre_id={}\n",
    "pre_best={}\n",
    "def get_true(city,Id):\n",
    "    infection=infection_data[city]\n",
    "    id_infect=infection[infection['region_id'].isin([Id])].reset_index(drop=True)\n",
    "    return id_infect['index'].values\n",
    "grouped=re_sub.groupby(['city'])\n",
    "for city,group in grouped:\n",
    "    data=group.reset_index(drop=True)\n",
    "    grouped1=data.groupby(['date'])\n",
    "    pre_city=grouped1['index'].agg(np.sum).reset_index(drop=True)#30天每天新增人数\n",
    "    pre_best[city]=pre_city\n",
    "'''\n",
    "\n",
    "\n",
    "#for key in city_infect.keys():\n",
    "#    city_infect[key]=pd.concat((city_infect[key],pre_best[key]),axis=0).reset_index(drop=True)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SEIR模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def adj_lamda(lamda_init,t,decay_rate):\n",
    "    # 衰减系数\n",
    "    #decay_rate= 0.02+decay_rate\n",
    "    decay_steps = 4\n",
    "    # 迭代轮数\n",
    "    global_steps = 90\n",
    "    # 指数感染率衰减过程\n",
    "    decayed_lamda = lamda_init * decay_rate**(t / decay_steps)\n",
    "    return decayed_lamda\n",
    "#定义误差\n",
    "def RMSE(data1,data2):\n",
    "    return np.sqrt(np.mean((data1-data2)**2))\n",
    "def RMSLE(data1,data2):\n",
    "    return np.sqrt(np.mean((np.log1p(data1)-np.log1p(data2))**2))\n",
    "#找到第一次出现病例的那天\n",
    "def find_zero(data):\n",
    "    for i in range(len(data)):\n",
    "        if data[i]>0:\n",
    "            return i\n",
    "#寻找参数\n",
    "def get_param(x,city):\n",
    "    x=x[find_zero(x):]\n",
    "    # 对城市总人口的估计\n",
    "    N_dic ={'A':5e6,'B':0.2e6,'C':2e6,'D':0.3e6,'E':0.15e6,\n",
    "            'F':5e6,'G':0.2e6,'H':0.2e6,'I':0.15e6,'J':0.15e6,'K':0.3e6}\n",
    "    lamda_list=np.linspace(0.1,1,20) #lamda 50个候选值\n",
    "    init_list=np.linspace(0.1,10,10) #初始感染人数,3个值\n",
    "    N_init=N_dic[city]  #先估计一下各个区人口总数\n",
    "    N_list=np.linspace(50000,N_init*5,50)#人口总数给100个值\n",
    "    #总共需要搜索的次数(50*3*100)\n",
    "    best_param={}\n",
    "    loss=1e5\n",
    "    #decay_rate_city={'A':0.935,'B':0.93,'C':0.955,'D':0.93,'E':0.93,\n",
    "            #'F':0.925,'G':0.93,'H':0.93,'I':0.975,'J':0.955,'K':0.97}#1.44\n",
    "    #decay_rate_city={'A':0.935,'B':0.96,'C':0.955,'D':0.96,'E':0.96,\n",
    "            #'F':0.95,'G':0.96,'H':0.97,'I':0.98,'J':0.96,'K':0.98}#1.33\n",
    "    #decay_rate_city={'A':0.93,'B':0.96,'C':0.98,'D':0.96,'E':0.96,\n",
    "            #'F':0.96,'G':0.95,'H':0.95,'I':0.985,'J':0.96,'K':0.975}#1.28\n",
    "    decay_rate_city={'A':0.93,'B':0.94,'C':0.98,'D':0.945,'E':0.955,\n",
    "            'F':0.944,'G':0.94,'H':0.955,'I':0.985,'J':0.96,'K':0.975}#1.04\n",
    "    for lamda_init in lamda_list:\n",
    "        for init_person in init_list:\n",
    "            for N in N_list:\n",
    "                gamma = 0.05 # recover rate\n",
    "                sigma = 1 / 3# exposed period\n",
    "                T=90\n",
    "                s = np.zeros([T])\n",
    "                e = np.zeros([T])\n",
    "                i = np.zeros([T])\n",
    "                r = np.zeros([T])\n",
    "                add=np.zeros([T])\n",
    "                i[0] = init_person/N #初始感染人数估计\n",
    "                add[0]=i[0]#\n",
    "                s[0] = 1\n",
    "                e[0] = init_person/N #初始潜伏期人数估计\n",
    "                for t in range(T-1):\n",
    "                    lamda=adj_lamda(lamda_init,t,decay_rate_city[city])\n",
    "                    s[t + 1] = s[t] - lamda * s[t] * i[t]\n",
    "                    e[t + 1] = e[t] + lamda * s[t] * i[t] - sigma * e[t]\n",
    "                    i[t + 1] = i[t] + sigma * e[t] - gamma * i[t]\n",
    "                    add[t+1] = sigma * e[t]\n",
    "                    r[t + 1] = r[t] + gamma * i[t]\n",
    "                tmp=RMSLE(add[60-len(x):60]*N,x)\n",
    "                if tmp<loss:\n",
    "                    loss=tmp\n",
    "                    best_param['lamda']=lamda_init\n",
    "                    best_param['N']=N\n",
    "                    best_param['init']=init_person\n",
    "                    best_param['value']=add*N\n",
    "                    \n",
    "    return best_param\n",
    "\n",
    "#\n",
    "def search_param():\n",
    "    pre_dic={}\n",
    "    for city in ['A','B','C','D','E','F','G','H','I','J','K']:\n",
    "        infection=city_infect[city]\n",
    "        y=infection.values\n",
    "        #10471, 7622, 95405, 12777, 3695\n",
    "        re=get_param(y,city)\n",
    "        pre_dic[city]=re['value'][60:]\n",
    "        print(\"city_%s的最佳拟合参数\"%city)\n",
    "        print(re['lamda'],re['N'],re['init'])\n",
    "        #画图检验搜索出来参数的真实拟合效果\n",
    "        fig, ax = plt.subplots(figsize=(10,6))\n",
    "        ax.plot(re['value'], c='r', lw=2, label='SEIR')\n",
    "        ax.plot(y, c='g', lw=2, label='true')\n",
    "        ax.set_xlabel('Day',fontsize=20)\n",
    "        ax.set_ylabel('Infect count', fontsize=20)\n",
    "        ax.grid(1)\n",
    "        plt.title(city,fontsize=20)\n",
    "        plt.xticks(fontsize=20)\n",
    "        plt.yticks(fontsize=20)\n",
    "        plt.legend();\n",
    "        save_dir='SEIR_fit/'\n",
    "        if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "        #plt.savefig(os.path.join(save_dir,city+'.png'))\n",
    "        plt.show()\n",
    "        plt.close()\n",
    "    return pre_dic\n",
    "#开始搜索参数\n",
    "pre_dicSeir=search_param()\n",
    "#\n",
    "#为每个区域分配新增人数，生成提交的csv文件\n",
    "import random\n",
    "import os\n",
    "#以最后一天为基准,求一个比例，给各个区域分配预测值\n",
    "def get_ratio(infection):\n",
    "    grouped=infection.groupby(['date'])\n",
    "    last_df=grouped.get_group(21200629).reset_index(drop=True)\n",
    "    ratio=last_df['index']/last_df['index'].sum()\n",
    "    return ratio.values\n",
    "ratios={}\n",
    "for city in ['A','B','C','D','E','F','G','H','I','J','K']:\n",
    "    ratios[city]=get_ratio(infection_data[city])\n",
    "#写结果\n",
    "pre_dic=pre_dicSeir#Seir的预测结果\n",
    "submit=pd.read_csv('submits/submission.csv',header=None)\n",
    "submit.columns = ['city','region_id','date', 'index']\n",
    "#\n",
    "submit=submit.drop(['index'],axis=1)\n",
    "pres=[]\n",
    "def convert_date(date):\n",
    "    if date==21200630:\n",
    "        return 0\n",
    "    else:\n",
    "        return date-21200700\n",
    "def get_result(city,region_id,date):\n",
    "    tmp=random.choice([0,0,2,2,1,2,2,1,2,0,0,1,0,2,0,2,2,1])\n",
    "    if city=='A':\n",
    "        pre=pre_dic['A'][convert_date(date)]*ratios['A'][region_id]+tmp\n",
    "    if city=='B':\n",
    "        pre=pre_dic['B'][convert_date(date)]*ratios['B'][region_id]+tmp\n",
    "    if city=='C':\n",
    "        pre=pre_dic['C'][convert_date(date)]*ratios['C'][region_id]+tmp\n",
    "    if city=='D':\n",
    "        pre=pre_dic['D'][convert_date(date)]*ratios['D'][region_id]+tmp\n",
    "    if city=='E':\n",
    "        pre=pre_dic['E'][convert_date(date)]*ratios['E'][region_id]+tmp\n",
    "    if city=='F':\n",
    "        pre=pre_dic['F'][convert_date(date)]*ratios['F'][region_id]+tmp\n",
    "    if city=='G':\n",
    "        pre=pre_dic['G'][convert_date(date)]*ratios['G'][region_id]+tmp\n",
    "    if city=='H':\n",
    "        pre=pre_dic['H'][convert_date(date)]*ratios['H'][region_id]+tmp\n",
    "    if city=='I':\n",
    "        pre=pre_dic['I'][convert_date(date)]*ratios['I'][region_id]+tmp\n",
    "    if city=='J':\n",
    "        pre=pre_dic['J'][convert_date(date)]*ratios['J'][region_id]+tmp\n",
    "    if city=='K':\n",
    "        pre=pre_dic['K'][convert_date(date)]*ratios['K'][region_id]+tmp\n",
    "    return pre\n",
    "#\n",
    "for _,city,region_id,date in submit.itertuples():\n",
    "    pre=get_result(city,region_id,date)\n",
    "    if pre<0:\n",
    "        pre=0\n",
    "    pres.append(int(pre))\n",
    "\n",
    "submit['index']=pres\n",
    "submit_dir='submits/'\n",
    "if not os.path.exists(submit_dir):\n",
    "    os.makedirs(submit_dir)\n",
    "submit.to_csv(os.path.join(submit_dir,'submission7_31v1.csv'),index=False,header=False)\n",
    "#\n",
    "def RMSLE(data1,data2):\n",
    "    return np.sqrt(np.mean((np.log1p(data1)-np.log1p(data2))**2))\n",
    "re_sub1=pd.read_csv('submits/submissionShilft1_16.csv',header=None)\n",
    "re_sub1.columns = ['city','region_id','date', 'index']\n",
    "re_sub2=pd.read_csv('submits/submission7_31v1.csv',header=None)\n",
    "re_sub2.columns = ['city','region_id','date', 'index']\n",
    "RMSLE(re_sub1['index'],re_sub2['index'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qyl",
   "language": "python",
   "name": "qyl"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
